#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <chrono>
#include "defs.h"

extern "C"
{

    void MY_MMult(int m, int n, int k, double *a, int lda,
                  double *b, int ldb,
                  double *c, int ldc)
    {

        size_t size = m * n * sizeof(double);

        // 在设备上分配内存
        double *d_A, *d_B, *d_C;
        cudaMalloc((void **)&d_A, size);
        cudaMalloc((void **)&d_B, size);
        cudaMalloc((void **)&d_C, size);

        // 将数据从主机复制到设备
        cudaMemcpy(d_A, a, size, cudaMemcpyHostToDevice);
        cudaMemcpy(d_B, b, size, cudaMemcpyHostToDevice);

        // 创建cuBLAS句柄
        cublasHandle_t handle;
        cublasCreate(&handle);

        double alpha = 1.0;
        double beta = 1.0;

        cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                    m, n, k,
                    &alpha, d_A, n,
                    d_B, n,
                    &beta, d_C, n);

        cudaMemcpy(c, d_C, size, cudaMemcpyDeviceToHost);

        cudaFree(d_A);
        cudaFree(d_B);
        cudaFree(d_C);
        cublasDestroy(handle);
    }
}
